"""
Calculous模块主要包括微积分, 常微分方程等相应的方法.
"""
import numpy as np


def derivative(f, self_value, level=1, self_varible_name=None, other_varibles=None, dx=0.01):
    """
    Derivitive
    ===
    差分方法求函数在某一点处的导数或偏导数.
    f为函数.
    self_value为所求导数点的自变量值.
    dx为差值.
    self_varible_name为求偏导数时的对象, 为字符串, 缺省则求导数.
    other_varibles为其他变量的取值, 为字典的形式. 如:
        dic = {'y':3.2, 'z':0.6}
    为y取3.2, z取0.6时对x的偏导数
    """
    if level == 1:
        if other_varibles == None or self_varible_name == None:
            return (f(self_value + dx) - f(self_value)) / dx
        dic = {}
        for key in f.__code__.co_varnames:
            if key == self_varible_name:
                dic[self_varible_name] = self_value
                continue
            if other_varibles.get(key) != None:
                dic[key] = other_varibles[key]
        v1 = f(**dic)
        for key in f.__code__.co_varnames:
            if key == self_varible_name:
                dic[self_varible_name] = self_value + dx
                break
        v2 = f(**dic)
        return (v2 - v1) / dx
    if level < 1 or type(level) != int:
        raise TypeError('不合法的导数阶数:{}'.format(level))
    return (derivative(f, self_value + dx, level-1, self_varible_name, other_varibles, dx) - derivative(f, self_value, level - 1, self_varible_name, other_varibles, dx)) / dx


def integral(f, lower, upper, self_varible_name=None, other_varibles=None, dx=0.01):
    """
    Derivitive
    ===
    差分方法求函数的定积分.
    f为函数.
    lower为积分下限. 
    upper为积分上限.
    self_varible_name为自变量名称. 缺省则为单变量函数.
    other_varibles为其他变量取值. 缺省则为单变量函数.
    dx为差值, 缺省则为0.01.
    """
    ans = 0
    for x in np.arange(lower, upper, dx):
        if other_varibles == None or self_varible_name == None:
            ans += f(x) * dx
        else:
            dic = {}
            for key in f.__code__.co_varnames:
                if key == self_varible_name:
                    dic[self_varible_name] = x
                    continue
                if other_varibles.get(key) != None:
                    dic[key] = other_varibles[key]
            ans += f(**dic) * dx
    return ans


def ode(f, level, self_final_value, target_init_values, self_init_value=0, self_varible_name=None, target_varible_name=None, output_level=None, dt=0.01):
    """
    Ordinary Derivative Equation
    ===
    使用有限元差分的方法求解偏微分方程的数值解.

    f: 常微分方程, 其参数是因变量的导数向量和自变量. 若不输入自变量和因变量, 默认因变量在自变量之前
    例如在四阶微分方程
        def function(y, x):
            return 2 * y[3] + np.sin(y[2]) - 1 / y[1] + y[0] ** 2 + 3
    此函数代表微分方程
        d^4y/dx^4 = 2 * d^3y/dx^3 + sin(d^2y/dx^2 - 1 / dy/dx + y ^ 2 + 3)

    level:常微分方程的阶数.
    self_final_value: 自变量的最终值.
    target_init_values: 因变量所有阶导数的初始值的向量.
    例如[0, 7.2, 3, 4]代表因变量的初始值为0, 因变量一阶导数的初始值为7.2, 二阶导数初始值为3, 三阶导数初始值为4.

    self_init_value: 自变量的初始值, 默认为0
    self_varible_name: 自变量的名称(字符串).
    target_varible_name: 因变量的名称(字符串).
    output_level: 输出的导数阶数, 缺省则输出整个导数向量.
    dt代表每一次计算的自变量差值.
    """
    dirseq = target_init_values  # 导数序列, 从
    ans = []
    dic = {}
    for t in np.arange(self_init_value, self_final_value, dt):
        if self_varible_name == None or target_varible_name == None:
            dirseq[level] = f(dirseq, t)
        else:
            for key in f.__code__.co_varnames:
                if self_varible_name == key:
                    dic[key] = t
                if target_varible_name == key:
                    dic[key] = dirseq
            dirseq[level] = f(**dic)
        for i in range(level - 1, -1, -1):
            a = dirseq[i] + dirseq[i + 1] * dt
            dirseq[i] = a
        if output_level == None:
            ans.append(tuple(dirseq))
        else:
            ans.append(dirseq[output_level])
    return ans
